# 16124278 王浩 流浪地球 Run_Cnn
# -*- coding: utf-8 -*-

from __future__ import print_function

import os
import time  # 时间
from datetime import timedelta

import numpy as np
import tensorflow as tf
from sklearn import metrics

from Cnn_Model import TCNNConfig, TextCNN  # 自己配置的CnnModel
from comments.comments_loader import read_vocab, read_category, batch_iter, process_file, build_vocab  # 对影评数据的操作

base_dir = './comments'
train_dir = os.path.join(base_dir, 'comments.train.txt')  # 训练集
test_dir = os.path.join(base_dir, 'comments.test.txt')  # 测试集
val_dir = os.path.join(base_dir, 'comments.val.txt')  # 验证集
vocab_dir = os.path.join(base_dir, 'comments.vocab.txt')  # 影评字典

save_dir = './checkpoints/textcnn'
save_path = os.path.join(save_dir, 'best_validation')  # 最佳验证结果保存路径



# 获取已使用时间
def get_time_dif(start_time):
    end_time = time.time()
    time_dif = end_time - start_time
    return timedelta(seconds=int(round(time_dif)))


# 添加参数
def feed_data(x_batch, y_batch, keep_prob):
    feed_dict = {
        model.input_x: x_batch,
        model.input_y: y_batch,
        model.keep_prob: keep_prob
    }
    return feed_dict


# 评估在某一数据上的准确率和损失
def evaluate(sess, x_, y_):
    data_len = len(x_)
    batch_eval = batch_iter(x_, y_, 128)
    total_loss = 0.0
    total_acc = 0.0
    for x_batch, y_batch in batch_eval:
        batch_len = len(x_batch)
        feed_dict = feed_data(x_batch, y_batch, 1.0)
        loss, acc = sess.run([model.loss, model.acc], feed_dict=feed_dict)
        total_loss += loss * batch_len
        total_acc += acc * batch_len

    return total_loss / data_len, total_acc / data_len


def train():
    print("正在配置TensorBoard和Saver.....")  # 这个是可视化的参数保存处，也就是每次训练的时候可以在这里看参数的变化
    # 配置 Tensorboard，重新训练时，需要将tensorboard文件夹删除，不然图会覆盖
    tensorboard_dir = './tensorboard/textcnn'
    if not os.path.exists(tensorboard_dir):
        os.makedirs(tensorboard_dir)

    tf.summary.scalar("loss", model.loss)  # 记录loss值
    tf.summary.scalar("accuracy", model.acc)  # 记录准确率
    merged_summary = tf.summary.merge_all()
    writer = tf.summary.FileWriter(tensorboard_dir)  # 记录写入

    # 配置 Saver
    saver = tf.train.Saver()  # 即checkpoint存放处,这个是参数存放处，可以继续训练或者保存最好的模型
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    print("载入训练集与验证集......")
    # 载入训练集与验证集
    start_time = time.time()  # 记录开始时间

    x_train, y_train = process_file(train_dir, word_to_id, cat_to_id, config.seq_length)  # 训练集数据预处理
    x_val, y_val = process_file(val_dir, word_to_id, cat_to_id, config.seq_length)  # 验证集数据预处理
    time_dif = get_time_dif(start_time)
    print("用时:", time_dif)

    # 创建session
    session = tf.Session()
    session.run(tf.global_variables_initializer())
    writer.add_graph(session.graph)  # 写入graph到session

    print('正在训练评估中...')
    start_time = time.time()  # 记录开始时间
    total_batch = 0  # 总批次
    best_acc_val = 0.0  # 最佳验证集准确率
    last_improved = 0  # 记录上一次提升批次
    require_improvement = 600  # 如果超过600轮未提升，提前结束训练

    flag = False
    for epoch in range(config.num_epochs):
        print('Epoch:', epoch + 1)
        batch_train = batch_iter(x_train, y_train, config.batch_size)  # 按batch_size切换批次
        for x_batch, y_batch in batch_train:
            feed_dict = feed_data(x_batch, y_batch, config.dropout_keep_prob)  # 将三个数据和标签放在一块，是model的传参
            # print("x_batch is {}".format(x_batch.shape))

            if total_batch % config.save_per_batch == 0:
                # 每多少轮次将训练结果写入tensorboard scalar
                s = session.run(merged_summary, feed_dict=feed_dict)
                writer.add_summary(s, total_batch)

            if total_batch % config.print_per_batch == 0:
                # 每多少轮次输出在训练集和验证集上的性能
                feed_dict[model.keep_prob] = 1.0
                loss_train, acc_train = session.run([model.loss, model.acc], feed_dict=feed_dict)
                loss_val, acc_val = evaluate(session, x_val, y_val)  # todo

                if acc_val > best_acc_val:
                    # 保存最好结果
                    best_acc_val = acc_val
                    last_improved = total_batch  # 保存
                    saver.save(sess=session, save_path=save_path)
                    improved_str = '*'  # 最好的结果展示
                else:
                    improved_str = ''

                time_dif = get_time_dif(start_time)
                msg = 'Iter: {0:>6}, Train Loss: {1:>6.2}, Train Acc: {2:>7.2%},' \
                      + ' Val Loss: {3:>6.2}, Val Acc: {4:>7.2%}, Time: {5} {6}'
                print(
                    msg.format(total_batch, loss_train, acc_train, loss_val, acc_val, time_dif, improved_str))  # 打印输出参数

            session.run(model.optim, feed_dict=feed_dict)  # 运行优化 真正开始运行,因为是相互依赖，倒着找的
            total_batch += 1

            if total_batch - last_improved > require_improvement or acc_val > 0.98:
                # 验证集正确率长期不提升，提前结束训练
                print("验证集正确率长期不提升，自动提前结束训练！！！")
                flag = True
                break  # 跳出循环
        if flag:  # 同上
            break


# 测试
def test():
    print("加载测试集......")
    start_time = time.time()
    x_test, y_test = process_file(test_dir, word_to_id, cat_to_id, config.seq_length)

    # 创建一个session
    session = tf.Session()
    session.run(tf.global_variables_initializer())
    saver = tf.train.Saver()
    saver.restore(sess=session, save_path=save_path)  # 读取保存的模型

    print('测试中......')
    loss_test, acc_test = evaluate(session, x_test, y_test)
    msg = 'Test Loss: {0:>6.2}, Test Acc: {1:>7.2%}'
    print(msg.format(loss_test, acc_test))  # 输出测试结果

    batch_size = 128  # 批次大小
    data_len = len(x_test)
    num_batch = int((data_len - 1) / batch_size) + 1  # 分多少次拿数据

    y_test_cls = np.argmax(y_test, 1)
    y_pred_cls = np.zeros(shape=len(x_test), dtype=np.int32)  # 保存预测结果
    for i in range(num_batch):  # 逐批次处理
        start_id = i * batch_size
        end_id = min((i + 1) * batch_size, data_len)
        feed_dict = {
            model.input_x: x_test[start_id:end_id],
            model.keep_prob: 1.0
        }  # 数据传入
        y_pred_cls[start_id:end_id] = session.run(model.y_pred_cls, feed_dict=feed_dict)

    # 评估
    print("准确率，召回率，F1均值.....")
    print(metrics.classification_report(y_test_cls, y_pred_cls, target_names=categories))

    # 混淆矩阵
    print("混淆矩阵处理中......   ")
    cm = metrics.confusion_matrix(y_test_cls, y_pred_cls)
    print(cm)

    time_dif = get_time_dif(start_time)
    print("用时:", time_dif)


if __name__ == '__main__':

    config = TCNNConfig()  # 1.加载配置参数; 初始化右边的类之后得到左边的对象config
    if not os.path.exists(vocab_dir):  # 如果不存在词汇表，重建
        build_vocab(train_dir, vocab_dir, config.vocab_size)
    categories, cat_to_id = read_category()  # 制作分类目录
    words, word_to_id = read_vocab(vocab_dir)  # 建立词汇目录
    config.vocab_size = len(words)  # 5000个
    model = TextCNN(config)  # 初始化CNN模型
    option = 'test'  # 选择模式
    if option == 'train':
        train()
    else:
        test()
